//BNB mpi
//$Id: bnb_mpi.h 709 2012-05-28 15:13:33Z Oleg.Bulychov@gmail.com $
#pragma once

//__________________________________________________________________________________
#include "src/heo/include/easy_mpi.h"
#include "src/heo/include/stream_visitor.h"
#include "solvers/bnb/bnb_common.h"
#include "solvers/bnb/bnb_balancer_random.h"
//__________________________________________________________________________________
//__________________________________________________________________________________
template<class TCostType>
struct bnb_gather_data
{
    TCostType   best_ub_;
    int         stop_;
    size_t      pq_size_;
    int         dst_;

    bnb_gather_data()
    {
        accept(meta::set_default_values());
    }

    //static char const* get_Name() { return "bnb_gathered_data"; }
    //static char const* get_Section() { return get_Name(); }

    template<class V>
    void accept(V& v)
    {
        v(VAR_(best_ub),    TCostType());
        v(VAR_(stop),       0);
        v(VAR_(pq_size),    0);
        v(VAR_(dst),        0);
    }
};
//__________________________________________________________________________________
//__________________________________________________________________________________
template<class T>
class BNB_MPI:
    public BNB<T>
{
public:
    typedef typename BNB<T>::TEnvironment TEnvironment;
    typedef typename BNB<T>::TSolution TSolution;

    typedef typename BNB<T>::TDirection TDirection;
    typedef typename BNB<T>::TCostType  TCostType;

    typedef typename BNB<T>::TTask TTask;
    typedef typename BNB<T>::TBoundPredicate TBoundPredicate;

    typedef typename TEnvironment::TSolution TInOutSolution;
    typedef typename TEnvironment::TStatistics TStatistics;

    typedef bnb_gather_data<TCostType> bnb_gather;
public:
    using BNB<T>::rank_;
    using BNB<T>::env_;
    using BNB<T>::rnd_;
    using BNB<T>::first_ub_;
    using BNB<T>::best_ub_;
    using BNB<T>::current_;
    using BNB<T>::pq_;
    using BNB<T>::predictor_;
    using BNB<T>::max_step_;
    using BNB<T>::BNB_CONTINUE;
    using BNB<T>::BNB_PQ_EMPTY;
    using BNB<T>::BNB_STOP;
    using BNB<T>::BNB_ERROR;
    using BNB<T>::BNB_IDLE;

protected:
    bnb_balancer<mpi::recv, mpi::ibsend> balancer_;
    size_t spiral_init_size_;
    int max_init_step_;

public:
    BNB_MPI():
        spiral_init_size_(),
        max_init_step_()
    {}

    template<class V>
    void accept(V& v)
    {
        v(static_cast<BNB<T>&>(*this), "/");
        v(VAR_(spiral_init_size),  8);
        v(VAR_(max_init_step),     1000);
        v(balancer_);
    }

    void solve(Solution_Pool<TInOutSolution>& pool,
        TStatistics& stat, mpi::communicator& comm)
    {
        int np = comm.size();
        if (1 == np)
        {//:don't bother mpi
            BNB<T>::solve(pool, stat);
            return;
        }
        solve_MPI(pool, stat, comm, np);
    }

    void solve_MPI(Solution_Pool<TInOutSolution>& pool,
        TStatistics& stat,
        mpi::communicator& comm,
        int np)
    {
        ASSERT_(this->rank_ == comm.rank());
        ASSERT_(np > 1);//:see above

        this->init_From_Pool(pool);
        this->calc_Max_Step();

        predictor_.init(this->env_, this->rnd_);

        int pq_init = 0;

        stat.update(0, this->rank_, *this, this->env_);

        int step = 1;
        if (0 == this->rank_)
        {//:root initialization
            if (this->get_Debug_Level())
            {
#pragma omp critical(debug_output)
                std::cout << std::setw(2)
                    << rank_ << ": begin BNB initialization"
                    << lf;
            }
            this->first_ub_ = this->best_ub_ = current_.get_Ub(this->env_);
            predictor_.pre_Calculate(this->env_, this->rnd_);

            ref_ptr<TTask> root;
            root.attach(new TTask);
            root->init(env_, rnd_);
            pq_.push(root);

            for (;;)
            {
                int r = this->do_Step(step);
                stat.update(step, this->rank_, *this, this->env_);
                if (r != (int)BNB_CONTINUE)
                {
                    this->show_Status(r, step);
                    break;
                }
                if (pq_.size() >= np * spiral_init_size_)
                {
                    pq_init = 1;
                    break;
                }
                //:max_step_ can be too big for initialization
                //:if pq_ won't grow
                if (++step > std::min(max_init_step_, max_step_))
                {
                    this->show_Status(BNB_STOP, step);
                    break;
                }
            }

            comm.barrier();
            mpi::bcast(0, comm) << pq_init;

            //:spiral initialization
            if (pq_init)
            {
                mpi::bcast(0, comm) << step
                    << this->first_ub_
                    << this->best_ub_
                    << this->predictor_;

                int dst = 1;
                for (size_t i = 1; i < pq_.size(); ++dst)
                {
                    if (dst % np == 0)
                    {
                        ++i;
                        continue;
                    }
                    mpi::send ds(dst % np, comm);
                    this->send_Task(i, ds);
                    pq_.erase(i);
                }

                if (this->get_Debug_Level())
                {
#pragma omp critical(debug_output)
                    std::cout << std::setw(2) << this->rank_
                        << ":" << std::setw(5) << step
                        << ": BNB initialization complete"
                        << lf;
                }
            }

        }
        else
        {
            comm.barrier();
            mpi::bcast(0, comm) >> pq_init;

            if (pq_init)
            {
                mpi::bcast(0, comm) >> step
                    >> this->first_ub_
                    >> this->best_ub_
                    >> this->predictor_;
                //:recv tasks
                mpi::recv in(0, comm);
                for (size_t i = 0; i < this->spiral_init_size_; ++i)
                {
                    this->recv_Task(in);
                }
            }
        }

        if (pq_init)
        {
            int res = main_Loop(comm, stat, np, step);
            this->show_Status(res, step);
        }

        //?if (this->get_Debug_Level())
        {
#pragma omp critical(debug_output)
            {
                std::cout << std::setw(2) << this->rank_ << ": ";
                this->stat_.dump(std::cout);
            }
        }

        clear_MPI(comm);

        update_Pool(pool, comm, np);
    }

private:
    int main_Loop(mpi::communicator& comm, TStatistics& stat, int np, int& step)
    {
        int coop_rate = this->balancer_.get_Rate();
        int pending_stop = BNB_CONTINUE;

        Array<bnb_gather> all_data;
        all_data.resize(np);
        for (int i = 0; i < np; ++i)
            all_data[i].dst_ = i;    
        bnb_gather data;

        for (; step <= this->max_step_; ++step)
        {
            int r = this->do_Step(step);
            stat.update(step, this->rank_, *this, this->env_);
            switch(r)
            {
            case BNB_ERROR:
            case BNB_STOP:
                if (!coop_rate)
                    return r;
                pending_stop = r;
                break;
            case BNB_PQ_EMPTY:
                if (!coop_rate)
                    return r;
                if (this->params_.debug_out_rate_ && ((this->stat_.idle_ % this->params_.debug_out_rate_) == 0))
                    show_Status(BNB_IDLE, step);
                ++this->stat_.idle_;
                break;
            default:
                break;
            }
            if (!coop_rate)
                continue;

            //if (balancer_.is_Async_Mode())
            //TODO fix me...

            //:sync only yet
            if ((step % coop_rate) == 0)
            {
                data.best_ub_ = this->best_ub_;
                data.pq_size_ = this->pq_.size();
                data.stop_ = pending_stop;
                data.dst_ = all_data[this->rank_].dst_;

                comm.all_gather(&data, 1, &all_data[0]);

                int r = process_Data(&all_data[0], np);
                if (r != BNB_CONTINUE)
                    return r;

                this->balancer_.balance_Sync(*this, this->rank_, np, step,
                    comm, &all_data[0]);
                comm.test_Isend();
            }
        }
        return BNB_STOP;
    }

    int process_Data(bnb_gather* p, int np)
    {
        size_t max_size = 0;
        for (int i = 0; i < np; ++i, ++p)
        {
            bnb_gather& data = *p;
            if (data.stop_ != BNB_CONTINUE)
                return data.stop_;
            if (TBoundPredicate()(data.best_ub_, this->best_ub_))
                this->best_ub_ = data.best_ub_;
            if (data.pq_size_ > max_size)
                max_size = data.pq_size_;
        }
        if (!max_size)
            return BNB_PQ_EMPTY;
        return BNB_CONTINUE;
    }

    void update_Pool(Solution_Pool<TInOutSolution>& pool,
        mpi::communicator& comm,
        int np)
    {
        TSolution& final_sol = this->get_Final_Solution();
        if (0 == this->rank_)
        {
            pool.sol_list_.resize(np);
            pool.sol_list_[0].init(this->env_, this->rnd_);
            pool.sol_list_[0].assign(final_sol);
            if (this->env_->debug_level_)
                std::cout << "0: final cost = " << final_sol.get_Cost() << lf;
            TSolution tmp;
            tmp.init(this->env_, this->rnd_);
            for (int i = 1; i < np; ++i)
            {
                mpi::recv(i, comm) >> tmp;
                pool.sol_list_[i].init(this->env_, this->rnd_);
                pool.sol_list_[i].assign(tmp);
                if (this->env_->debug_level_)
                    std::cout << i << ": final cost = " << tmp.get_Cost() << lf;
            }
        }
        else
        {
            mpi::send(0, comm) << final_sol;
        }
        comm.barrier();
    }

    void clear_MPI(mpi::communicator& comm)
    {
        comm.barrier();
        comm.test_Isend();
        //if (balancer_.is_Async_Mode())
        //TODO fix me...
    }
};
//__________________________________________________________________________________
//__________________________________________________________________________________
//__________________________________________________________________________________
//__________________________________________________________________________________
//__________________________________________________________________________________

//EOF!
